#include "scan_ground_filter.hpp"

#include <pcl/filters/extract_indices.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl_conversions/pcl_conversions.h>

#include <cmath>
#include <algorithm>
#include <limits>

namespace autoware_ground_segmentation
{

ScanGroundFilter::ScanGroundFilter(ros::NodeHandle& nh, ros::NodeHandle& private_nh)
: nh_(nh), private_nh_(private_nh), tf2_buffer_(ros::Duration(10)), tf2_listener_(tf2_buffer_)
{
  // 获取参数
  private_nh_.param<bool>("elevation_grid_mode", elevation_grid_mode_, false);
  private_nh_.param<float>("non_ground_height_threshold", non_ground_height_threshold_, 0.15);
  private_nh_.param<float>("low_priority_region_x", low_priority_region_x_, 2.0);
  private_nh_.param<float>("center_pcl_shift", center_pcl_shift_, 0.0);
  
  private_nh_.param<float>("radial_divider_angle", radial_divider_angle_rad_, 0.1);
  private_nh_.param<float>("global_slope_max_angle", global_slope_max_angle_rad_, 0.2);
  private_nh_.param<float>("local_slope_max_angle", local_slope_max_angle_rad_, 0.1);
  private_nh_.param<float>("global_slope_max_ratio", global_slope_max_ratio_, 0.2);
  private_nh_.param<float>("local_slope_max_ratio", local_slope_max_ratio_, 0.1);
  private_nh_.param<float>("split_points_distance_tolerance", split_points_distance_tolerance_, 0.2);
  
  private_nh_.param<bool>("use_virtual_ground_point", use_virtual_ground_point_, false);
  private_nh_.param<float>("split_height_distance", split_height_distance_, 0.1);
  
  private_nh_.param<bool>("use_recheck_ground_cluster", use_recheck_ground_cluster_, false);
  private_nh_.param<float>("recheck_start_distance", recheck_start_distance_, 5.0);
  private_nh_.param<bool>("use_lowest_point", use_lowest_point_, false);
  private_nh_.param<float>("detection_range_z_max", detection_range_z_max_, 2.0);
  
  private_nh_.param<float>("grid_size_m", grid_size_m_, 0.1);
  private_nh_.param<float>("grid_mode_switch_radius", grid_mode_switch_radius_, 20.0);
  int temp_buffer_size;
  private_nh_.param<int>("gnd_grid_buffer_size", temp_buffer_size, 4);
  gnd_grid_buffer_size_ = static_cast<uint16_t>(temp_buffer_size);
  private_nh_.param<float>("virtual_lidar_z", virtual_lidar_z_, 2.0);

  // 计算派生参数
  radial_dividers_num_ = static_cast<size_t>(2.0 * M_PI / radial_divider_angle_rad_);
  
  // 设置PCL过滤器
  voxel_grid_.setLeafSize(0.1, 0.1, 0.1);

  // 设置ROS
  pointcloud_sub_ = nh_.subscribe("input", 1, &ScanGroundFilter::pointcloudCallback, this);
  ground_cloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("output/ground", 1);
  no_ground_cloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("output/no_ground", 1);

  ROS_INFO("Scan Ground Filter initialized");
}

void ScanGroundFilter::pointcloudCallback(const sensor_msgs::PointCloud2::ConstPtr& input_msg)
{
  // 转换为PCL格式
  pcl::PointCloud<pcl::PointXYZ>::Ptr input_cloud(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::fromROSMsg(*input_msg, *input_cloud);

  if (input_cloud->empty()) {
    ROS_WARN("Empty input point cloud");
    return;
  }

  // 体素网格过滤
  pcl::PointCloud<pcl::PointXYZ>::Ptr voxel_cloud(new pcl::PointCloud<pcl::PointXYZ>);
  voxel_grid_.setInputCloud(input_cloud);
  voxel_grid_.filter(*voxel_cloud);

  // 转换为有序格式
  std::vector<PointCloudVector> radial_ordered_points;
  convertPointcloud(input_msg, radial_ordered_points);

  // 分类点云
  pcl::PointIndices no_ground_indices;
  classifyPointCloud(input_msg, radial_ordered_points, no_ground_indices);

  // 输出分类结果
  ROS_INFO("Final result: total %zu, ground %zu, non-ground %zu", 
           input_cloud->size(), 
           input_cloud->size() - no_ground_indices.indices.size(),
           no_ground_indices.indices.size());
  
  // 提取并发布结果
  sensor_msgs::PointCloud2 ground_cloud, no_ground_cloud;
  
  // 创建地面点云（排除非地面点）
  pcl::PointCloud<pcl::PointXYZ>::Ptr ground_pcl(new pcl::PointCloud<pcl::PointXYZ>);
  ground_pcl->points.reserve(input_cloud->size() - no_ground_indices.indices.size());
  
  // 创建非地面点云（只包含非地面点）
  pcl::PointCloud<pcl::PointXYZ>::Ptr no_ground_pcl(new pcl::PointCloud<pcl::PointXYZ>);
  no_ground_pcl->points.reserve(no_ground_indices.indices.size());
  
  // 使用布尔数组标记非地面点
  std::vector<bool> is_no_ground(input_cloud->size(), false);
  
  for (const auto& idx : no_ground_indices.indices) {
    if (idx < input_cloud->size()) {
      is_no_ground[idx] = true;
    }
  }
  
  // 分离地面和非地面点
  for (size_t i = 0; i < input_cloud->size(); ++i) {
    if (is_no_ground[i]) {
      // 非地面点
      no_ground_pcl->points.push_back(input_cloud->points[i]);
    } else {
      // 地面点
      ground_pcl->points.push_back(input_cloud->points[i]);
    }
  }
  
  // 添加调试信息：验证点云分离
  ROS_INFO("Point clouds: ground %zu, non-ground %zu", 
           ground_pcl->size(), no_ground_pcl->size());
  
  // 创建带强度的点云
  pcl::PointCloud<pcl::PointXYZI>::Ptr ground_cloud_intensity(new pcl::PointCloud<pcl::PointXYZI>);
  pcl::PointCloud<pcl::PointXYZI>::Ptr no_ground_cloud_intensity(new pcl::PointCloud<pcl::PointXYZI>);
  
  // 地面点云：强度设置为10
  ground_cloud_intensity->points.reserve(ground_pcl->size());
  for (const auto& point : ground_pcl->points) {
    pcl::PointXYZI intensity_point;
    intensity_point.x = point.x;
    intensity_point.y = point.y;
    intensity_point.z = point.z;
    intensity_point.intensity = 10.0;  // 地面点强度设置为10
    ground_cloud_intensity->points.push_back(intensity_point);
  }
  
  // 非地面点云：强度设置为30
  no_ground_cloud_intensity->points.reserve(no_ground_pcl->size());
  for (const auto& point : no_ground_pcl->points) {
    pcl::PointXYZI intensity_point;
    intensity_point.x = point.x;
    intensity_point.y = point.y;
    intensity_point.z = point.z;
    intensity_point.intensity = 30.0;  // 非地面点强度设置为30
    no_ground_cloud_intensity->points.push_back(intensity_point);
  }
  
  // 转换回ROS消息 - 使用带强度的点云
  pcl::toROSMsg(*ground_cloud_intensity, ground_cloud);
  pcl::toROSMsg(*no_ground_cloud_intensity, no_ground_cloud);
  
  ground_cloud.header = input_msg->header;
  no_ground_cloud.header = input_msg->header;
  
  // 发布结果
  ground_cloud_pub_.publish(ground_cloud);
  no_ground_cloud_pub_.publish(no_ground_cloud);
}

void ScanGroundFilter::convertPointcloud(
  const sensor_msgs::PointCloud2::ConstPtr & in_cloud,
  std::vector<PointCloudVector> & out_radial_ordered_points) const
{
  pcl::PointCloud<pcl::PointXYZ>::Ptr pcl_cloud(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::fromROSMsg(*in_cloud, *pcl_cloud);
  
  out_radial_ordered_points.clear();
  out_radial_ordered_points.resize(radial_dividers_num_);
  
  for (size_t i = 0; i < pcl_cloud->size(); ++i) {
    const auto& point = pcl_cloud->points[i];
    
    // 计算半径和角度
    float radius = std::sqrt(point.x * point.x + point.y * point.y);
    float theta = std::atan2(point.y, point.x);
    if (theta < 0) theta += 2.0 * M_PI;
    
    // 确定径向分割
    size_t radial_div = static_cast<size_t>(theta / radial_divider_angle_rad_);
    if (radial_div >= radial_dividers_num_) radial_div = radial_dividers_num_ - 1;
    
    // 创建点数据
    PointData point_data;
    point_data.radius = radius;
    point_data.data_index = i;
    point_data.grid_id = static_cast<uint16_t>(radius / grid_size_m_);
    
    out_radial_ordered_points[radial_div].push_back(point_data);
  }
  
  // 对每个径向分割按半径排序
  for (auto& radial_points : out_radial_ordered_points) {
    std::sort(radial_points.begin(), radial_points.end(),
              [](const PointData& a, const PointData& b) {
                return a.radius < b.radius;
              });
  }
}

void ScanGroundFilter::calcVirtualGroundOrigin(pcl::PointXYZ & point) const
{
  point.x = center_pcl_shift_;
  point.y = 0.0;
  point.z = virtual_lidar_z_;
}

void ScanGroundFilter::classifyPointCloud(
  const sensor_msgs::PointCloud2::ConstPtr& input_msg,
  const std::vector<PointCloudVector>& radial_ordered_points,
  pcl::PointIndices& no_ground_indices)
{
  // 转换输入点云
  pcl::PointCloud<pcl::PointXYZ>::Ptr input_cloud(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::fromROSMsg(*input_msg, *input_cloud);
  
  no_ground_indices.indices.clear();
  
  // 虚拟地面点（车辆中心）
  pcl::PointXYZ virtual_ground_point(0, 0, 0);
  
  // 为每个射线运行分类算法
  for (size_t i = 0; i < radial_ordered_points.size(); ++i) {
    float prev_gnd_radius = 0.0f;
    float prev_gnd_slope = 0.0f;
    PointLabel point_label_curr = PointLabel::INIT;
    
    pcl::PointXYZ prev_gnd_point(0, 0, 0), point_curr, point_prev;
    
    // 遍历射线中的点
    for (size_t j = 0; j < radial_ordered_points[i].size(); ++j) {
      float points_distance = 0.0f;
      
      // 设置前一个点
      point_prev = point_curr;
      PointLabel point_label_prev = point_label_curr;
      
      // 设置当前点
      const PointData& pd = radial_ordered_points[i][j];
      point_label_curr = pd.point_state;
      
      // 获取当前点坐标
      point_curr = input_cloud->points[pd.data_index];
      
      if (j == 0) {
        // 第一个点
        if (use_virtual_ground_point_) {
          prev_gnd_point = virtual_ground_point;
        } else {
          prev_gnd_point = pcl::PointXYZ(point_curr.x, point_curr.y, point_curr.z);
        }
        prev_gnd_radius = std::hypot(prev_gnd_point.x, prev_gnd_point.y);
        prev_gnd_slope = 0.0f;
        
        // 第一个点默认设为地面点
        point_label_curr = PointLabel::GROUND;
        
        points_distance = std::sqrt(
          std::pow(point_curr.x - prev_gnd_point.x, 2) + 
          std::pow(point_curr.y - prev_gnd_point.y, 2) + 
          std::pow(point_curr.z - prev_gnd_point.z, 2));
      } else {
        // 计算与前一个点的距离
        points_distance = std::sqrt(
          std::pow(point_curr.x - point_prev.x, 2) + 
          std::pow(point_curr.y - point_prev.y, 2) + 
          std::pow(point_curr.z - point_prev.z, 2));
      }
      
      // 计算与地面点的距离和高度差
      float radius_distance_from_gnd = pd.radius - prev_gnd_radius;
      float height_from_gnd = point_curr.z - prev_gnd_point.z;
      
      // 检查点是否足够接近前一个点
      bool is_point_close_to_prev = (points_distance < 
        (pd.radius * radial_divider_angle_rad_ + split_points_distance_tolerance_));
      
      // 计算全局斜率比例
      float global_slope_ratio = pd.radius > 0.0f ? point_curr.z / pd.radius : 0.0f;
      
      bool calculate_slope = true;
      
      // 全局斜率检查
      if (global_slope_ratio > global_slope_max_ratio_) {
        point_label_curr = PointLabel::NON_GROUND;
        calculate_slope = false;
      } else if (point_label_prev == PointLabel::GROUND && is_point_close_to_prev &&
                 std::abs(height_from_gnd) < split_height_distance_) {
        // 接近前一个点，设置点跟随标签
        point_label_curr = PointLabel::POINT_FOLLOW;
        calculate_slope = false;
      }
      
      if (calculate_slope) {
        // 远离前一个点，计算局部斜率
        float local_slope = std::atan2(height_from_gnd, radius_distance_from_gnd);
        float slope_diff = local_slope - prev_gnd_slope;
        
        // 分类逻辑
        if (slope_diff > local_slope_max_angle_rad_) {
          // 点超出局部斜率阈值
          point_label_curr = PointLabel::NON_GROUND;
        } else {
          point_label_curr = PointLabel::GROUND;
        }
      }
      
      // 处理点跟随标签
      if (point_label_curr == PointLabel::POINT_FOLLOW) {
        point_label_curr = PointLabel::GROUND;
      }
      
      // 更新地面状态
      if (point_label_curr == PointLabel::GROUND) {
        // 正常更新参考点
        prev_gnd_radius = pd.radius;
        prev_gnd_point = pcl::PointXYZ(point_curr.x, point_curr.y, point_curr.z);
        prev_gnd_slope = std::atan2(height_from_gnd, radius_distance_from_gnd);
      }
      
      // 添加到非地面索引
      if (point_label_curr == PointLabel::NON_GROUND) {
        no_ground_indices.indices.push_back(pd.data_index);
      }
    }
  }
  
  // 输出分类结果
  ROS_INFO("Classification completed: total %zu, non-ground %zu", 
           input_cloud->size(), no_ground_indices.indices.size());
}

// 从点云中提取对象点
void ScanGroundFilter::extractObjectPoints(
  const sensor_msgs::PointCloud2::ConstPtr & in_cloud_ptr,
  const pcl::PointIndices & in_indices,
  sensor_msgs::PointCloud2 & out_object_cloud) const
{
  // 初始化输出点云
  out_object_cloud = *in_cloud_ptr;
  out_object_cloud.data.clear();
  out_object_cloud.width = 0;
  out_object_cloud.height = 0;
  
  if (in_indices.indices.empty()) {
    return;
  }
  
  // 转换为PCL格式进行处理
  pcl::PointCloud<pcl::PointXYZ>::Ptr input_pcl(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::fromROSMsg(*in_cloud_ptr, *input_pcl);
  
  // 使用PCL提取非地面点
  pcl::PointCloud<pcl::PointXYZ>::Ptr output_pcl(new pcl::PointCloud<pcl::PointXYZ>);
  
  // 创建临时extract_indices对象
  pcl::ExtractIndices<pcl::PointXYZ> temp_extract;
  temp_extract.setInputCloud(input_pcl);
  temp_extract.setIndices(boost::make_shared<pcl::PointIndices>(in_indices));
  temp_extract.setNegative(false);
  temp_extract.filter(*output_pcl);
  
  // 转换回ROS消息
  pcl::toROSMsg(*output_pcl, out_object_cloud);
  out_object_cloud.header = in_cloud_ptr->header;
}

} // namespace autoware_ground_segmentation

int main(int argc, char** argv)
{
  ros::init(argc, argv, "scan_ground_filter_node");
  
  ros::NodeHandle nh;
  ros::NodeHandle private_nh("~");
  
  autoware_ground_segmentation::ScanGroundFilter filter(nh, private_nh);
  
  ros::spin();
  
  return 0;
} 